# src/algorithms/rectification.py
import cv2
import numpy as np
from typing import Tuple
from src.data.structures import BinocularImagePair, CalibrationData


class BinocularRectifier:
    """双目视觉图像对的校正器

    该类用于对立体相机拍摄的图像对进行校正, 使得两幅图像的极线对齐, 便于后续的立体匹配
    """

    def __init__(self, calibration_data: CalibrationData):
        """初始化立体校正器.

        Args:
            calibration_data (CalibrationData): 相机标定数据, 包含相机内参和外参.
        """
        self._calib_data = calibration_data
        self._rectification_maps = None
        self._roi_left = None
        self._roi_right = None
        self._final_roi = None  # 用于存储最终裁剪区域
        self._Q = None

        self._computeRectificationMaps()

    def _computeRectificationMaps(self):
        """计算用于图像校正的映射表.

        基于相机标定数据计算立体校正所需的映射表. 该方法在类初始化时自动调用.
        主要步骤包括:
        1. 获取或推断图像尺寸
        2. 计算立体校正参数 (旋转矩阵和投影矩阵)
        3. 生成左右相机的重映射表
        """
        # 自动推断图像尺寸 (优先从投影矩阵或内参矩阵获取)
        # 若 CalibrationData 没有保存图像尺寸, 可在外部传入或在此处补充参数
        if (
            hasattr(self._calib_data, "image_size")
            and self._calib_data.image_size is not None
        ):
            image_size = self._calib_data.image_size
        else:
            # 尝试从内参矩阵获取 principal point 估算尺寸
            K = self._calib_data.camera_matrix_left
            cx, cy = int(K[0, 2] * 2), int(K[1, 2] * 2)
            image_size = (cx, cy)
        # image_size 格式为 (width, height)

        # 立体矫正
        # flag=CALIB_ZERO_DISPARITY: 强制让矫正后左右主点在同一行, 视差图第 0 行就是无穷远平面, 后续做立体匹配更稳
        # R1: 左相机的矫正旋转矩阵 (3x3)
        # R2: 右相机的矫正旋转矩阵 (3x3)
        # P1: 左相机的新投影矩阵 (3x4)
        # P2: 右相机的新投影矩阵 (3x4)
        # Q: 视差-深度重投影矩阵 (4x4)
        # roi_left: 左图像有效区域 (x, y, w, h)
        # roi_right: 右图像有效区域 (x, y, w, h)
        R1, R2, P1, P2, Q, roi_left, roi_right = cv2.stereoRectify(
            self._calib_data.camera_matrix_left,
            self._calib_data.dist_coeffs_left,
            self._calib_data.camera_matrix_right,
            self._calib_data.dist_coeffs_right,
            image_size,
            self._calib_data.rotation_matrix,
            self._calib_data.translation_vector,
            flags=cv2.CALIB_ZERO_DISPARITY,
        )

        # 计算矫正映射
        # map1 存目标像素坐标 (x,y) 的"源整数坐标"+"小数部分", map2 存插值权重
        # 注意: 映射表只与相机参数、图像分辨率有关, 与具体图像内容无关, 只需计算一次
        map1_left, map2_left = cv2.initUndistortRectifyMap(
            self._calib_data.camera_matrix_left,
            self._calib_data.dist_coeffs_left,
            R1,
            P1,
            image_size,
            cv2.CV_16SC2,
        )

        map1_right, map2_right = cv2.initUndistortRectifyMap(
            self._calib_data.camera_matrix_right,
            self._calib_data.dist_coeffs_right,
            R2,
            P2,
            image_size,
            cv2.CV_16SC2,
        )

        self._rectification_maps = (map1_left, map2_left, map1_right, map2_right)
        self._roi_left = roi_left
        self._roi_right = roi_right
        self._Q = Q

        print("Rectification maps computed successfully")

    def rectifyPair(
        self, stereo_pair: BinocularImagePair, crop_to_valid: bool = True
    ) -> BinocularImagePair:
        """对输入的立体图像对进行校正.

        对输入的立体图像对应用预计算的映射表进行校正, 使两幅图像的极线对齐.
        校正后的图像对可直接用于立体匹配算法.

        Args:
            stereo_pair (StereoImagePair): 待校正的立体图像对.
            crop_to_valid (bool): 是否裁剪到有效区域. 默认 True.

        Returns:
            StereoImagePair: 校正后的立体图像对.

        Raises:
            ValueError: 如果校正映射表未计算 (rectification_maps 为 None).
            ValueError: 如果裁剪后有效区域为空.
        """
        if self._rectification_maps is None:
            raise ValueError("Rectification maps not computed")

        map1_left, map2_left, map1_right, map2_right = self._rectification_maps

        # 应用矫正映射
        rectified_left = cv2.remap(
            stereo_pair.left_image, map1_left, map2_left, cv2.INTER_LINEAR
        )
        rectified_right = cv2.remap(
            stereo_pair.right_image, map1_right, map2_right, cv2.INTER_LINEAR
        )

        if crop_to_valid and self._roi_left is not None and self._roi_right is not None:
            # 计算左右图像共同的有效区域
            # 垂直方向: 取共同的有效范围
            y_start = max(self._roi_left[1], self._roi_right[1])
            y_end = min(
                self._roi_left[1] + self._roi_left[3],
                self._roi_right[1] + self._roi_right[3],
            )
            h = max(0, y_end - y_start)

            # 水平方向：计算有效重叠区域宽度
            # 为确保立体匹配的准确性, 左右图像必须保持相同宽度
            # 使用较小的有效宽度, 确保两幅图像裁剪后尺寸相同
            w = min(self._roi_left[2], self._roi_right[2])

            # 为保持对齐, 使用各自 ROI 的起始位置
            x_left = self._roi_left[0]
            x_right = self._roi_right[0]

            # 裁剪两幅图像到相同的尺寸
            if h > 0 and w > 0:
                rectified_left = rectified_left[
                    y_start : y_start + h, x_left : x_left + w
                ]
                rectified_right = rectified_right[
                    y_start : y_start + h, x_right : x_right + w
                ]
            else:
                raise ValueError(
                    "Valid region is empty, check calibration data or input images."
                )

        return BinocularImagePair(
            rectified_left, rectified_right, stereo_pair.timestamp
        )

    @property
    def reprojection_matrix(self) -> np.ndarray | None:
        """返回视差-深度重投影矩阵 Q.

        Returns:
            np.ndarray: 4x4 的重投影矩阵 Q.
        """
        return self._Q
